# Notice! As cicflowmeter (https://gitlab.com/juzai/cicflowmeter/-/tree/juzai-master-patch-32230)
# require numpy version no greater than 1.18, use flowing command to run this scripy:
# source cicflowmeter-venv/bin/activate && python extractor.py && deactivate

import subprocess, os, time
from multiprocessing import Pool

def extract(args):
    """Extact features from single pcap file.

    Args:
        pcap_file: pcap file path.
        csv_file: Features file path.
    """
    pcap_file, csv_file = args[0], args[1]
    cmd = ["cicflowmeter", "-f", pcap_file, "-c", csv_file]
    subprocess.call(cmd)

def main():
    """Extact features from pcap files.

    The structure of pcap_dir should be orgnaized as follow:

    dataset_path (path)
    ├── domain_1
    ├── domain_2
    ...
    └── domain_n

    where domain_i (e.g. apple.com) contains several pcap files.

    The structure of csv_dir will be like:
    csv_dir (path)
    └── time (filename)
        ├── domain_1
        ├── domain_2
        ...
        └── domain_n

    where domain_i (e.g. apple.com) contains several csv files.
    """
    config = {
        "dataset_path": "DoH_traffic_dataset",
        "csv_dir": "data/output_of_cicflowmeter"
    }

    dataset_path = config["dataset_path"]
    csv_dir = config["csv_dir"]

    domains = os.listdir(dataset_path)

    process_cap = 4

    with Pool(process_cap) as p:
        args = []

        csv_subdir_path = os.path.join(csv_dir, time.strftime("%Y%m%d%H%M%S", time.localtime()))
        pcap_subdir_path = dataset_path

        for domain in domains:
            pcap_domain_dir = os.path.join(pcap_subdir_path, domain)
            csv_domain_dir = os.path.join(csv_subdir_path, domain)

            if not os.path.exists(csv_domain_dir):
                os.makedirs(csv_domain_dir)

            pcap_file_names = os.listdir(pcap_domain_dir)
            for pcap_file_name in pcap_file_names:
                pcap_file_path = os.path.join(pcap_subdir_path, domain, pcap_file_name)
                csv_file_path = os.path.join(csv_subdir_path, domain, pcap_file_name[:-4] + "csv")
                args.append((pcap_file_path, csv_file_path))

        p.map(extract, args)

if __name__ == "__main__":
    t0 = time.time()
    main()
    t1 = time.time()

    print("Duration: ", t1 - t0)